local _G = _G

local error = _G.error
local pairs = _G.pairs
local setmetatable = _G.setmetatable
local unpack = _G.unpack
local math_sqrt = _G.math.sqrt

local CreateFrame = _G.CreateFrame
local GetLocale = _G.GetLocale
local UnitCanAttack = _G.UnitCanAttack
local UnitCreatureType = _G.UnitCreatureType
local UnitHealth = _G.UnitHealth
local UnitHealthMax = _G.UnitHealthMax
local UnitIsDead = _G.UnitIsDead
local UnitIsFriend = _G.UnitIsFriend
local UnitIsUnit = _G.UnitIsUnit
local UnitLevel = _G.UnitLevel
local UnitName = _G.UnitName
local UnitPlayerControlled = _G.UnitPlayerControlled

local AceLibrary = _G.AceLibrary

local MAJOR_VERSION = "MobHealth-1.0"
local MINOR_VERSION = "$Revision: 1 $"

if not AceLibrary then error(MAJOR_VERSION .. " requires AceLibrary") end
if not AceLibrary:IsNewVersion(MAJOR_VERSION, MINOR_VERSION) then return end
if not AceLibrary:HasInstance("AceEvent-2.0") then error(MAJOR_VERSION .. " requires AceEvent-2.0") end

local CREATURE_TYPE_DEMON, CREATURE_TYPE_BEAST
do
	local locale = GetLocale()
	if locale == "frFR" then
		CREATURE_TYPE_DEMON = "Démon"
		CREATURE_TYPE_BEAST = "Bête"
	elseif locale == "koKR" then
		CREATURE_TYPE_DEMON = "악마"
		CREATURE_TYPE_BEAST = "야수"
	elseif locale  == "zhCN" then
		CREATURE_TYPE_DEMON = "恶魔"
		CREATURE_TYPE_BEAST = "野兽"
	elseif locale  == "zhTW" then
		CREATURE_TYPE_DEMON = "惡魔"
		CREATURE_TYPE_BEAST = "野獸"
	elseif locale == "enUS" then
		CREATURE_TYPE_DEMON = "Demon"
		CREATURE_TYPE_BEAST = "Beast"
	end
end

local MobHealth = {}

local Debug = true

do
	local AceEvent = AceLibrary("AceEvent-2.0")
	AceEvent:embed(MobHealth)
end

local data_mt = {
	__mode = "k",
}

local damages = setmetatable({}, data_mt)
local heals = setmetatable({}, data_mt)
local healths = setmetatable({}, data_mt)

local unit_ids = {}
local db = {}

MobHealth.db = db

local focus_is_target
local get_unit_id = function (unit)
	if unit ~= "focus" or not focus_is_target then
		return unit_ids[unit]
	end
end

function MobHealth:OnEnable()
	self:RegisterEvent("UNIT_COMBAT")
	self:RegisterEvent("PLAYER_TARGET_CHANGED")
	self:RegisterEvent("PLAYER_FOCUS_CHANGED")
	self:RegisterEvent("UNIT_NAME_UPDATE", "UnitUpdate")
	self:RegisterEvent("UNIT_CLASSIFICATION_CHANGED", "UnitUpdate")
	self:RegisterEvent("UNIT_HEALTH")
end

function MobHealth:UNIT_COMBAT(unit, event, flags, amount)
	local id = get_unit_id(unit)
	if id then
		if event == "HEAL" then
			heals[id] = heals[id] + amount
		elseif event == "WOUND" then
			damages[id] = damages[id] + amount
		end
	end
end

function MobHealth:PLAYER_TARGET_CHANGED()
	self:UnitUpdate("player")
end

function MobHealth:PLAYER_FOCUS_CHANGED()
	self:UnitUpdate("focus")
end

function MobHealth:UNIT_HEALTH(unit)
	local id = get_unit_id(unit)
	if id then
		self:UpdateHealth(unit, id, UnitHealth(unit), UnitHealthMax(unit))
	end
end

local function set_lore_entry(id, healthmax)
	local entry = db[id]
	if not entry then
		entry = {
			level = "lore",
		}
		db[id] = entry
	else
		for k in pairs(entry) do
			entry[k] = nil
		end
		entry.level = "lore"
	end
	entry.hppp = healthmax / 100
end

function MobHealth:UnitUpdate(unit)
	local prev_id = unit_ids[unit]
	
	local id
	local type = UnitCreatureType(unit)
	if
		UnitCanAttack("player", unit) and
		not UnitIsDead(unit) and
		not UnitIsFriend("player", unit) and
		not (
			type == CREATURE_TYPE_DEMON or type == CREATURE_TYPE_BEAST and 
			UnitPlayerControlled(unit)
		)
	then
		id = UnitName(unit)..":"..UnitLevel(unit)
	end
	
	unit_ids[unit] = id
	
	if not id or id == prev_id then return end

	local health = UnitHealth(unit)
	local healthmax = UnitHealthMax(unit)
	
	if healthmax ~= 100 then
		set_lore_entry(id, healthmax)
		return
	end
	
	damages[id] = 0
	heals[id] = 0

	healths[id] = health
	
	focus_is_target = UnitIsUnit("focus", "target")
end

function MobHealth:UpdateHealth(unit, id, amount, healthmax)
	if healthmax ~= 100 then
		set_lore_entry(id, healthmax)
		return
	end
	
	local delta_percent = amount - healths[id]
	local delta_amount = heals[id] - damages[id]
	
	local entry = db[id]
	if not entry then
		entry = {
			level = "estimate",
		}
		db[id] = entry
	elseif entry.level ~= "estimate" then
		return
	end
	
	local percent = (amount - healths[id])
	local heals = heals[id]
	local damages = damages[id]
	
	if Debug then
		local samples = entry.samples
		if not samples then
			samples = {}
			entry.samples = samples
		end
		samples[#samples + 1] = { amount, percent, heals, damages }
		
		self:UpdateDebugInfo(unit, samples)
	end
	
	if percent >= 1 and heals == 0 then
		percent = percent + entry.percent
		damages = damages + entry.damages
	
		entry.percent = percent
		entry.damages = damages
		
		entry.hppp = percent / damages
	end
	
	damages[id] = 0
	heals[id] = 0
	healths[id] = amount
end

-- API
function MobHealth:GetUnitHealth(unit, amount, healthmax, name, level)
	local id = unit_ids[unit]
	if not id then
		id = (name or UnitName(unit))..":"..(level or UnitLevel(unit))
	end
	if not amount then amount = UnitHealth(unit) end
	if not healthmax then healthmax = UnitHealthMax(unit) end
	
	local entry = db[id]
	if not entry then
		return amount, healthmax
	end
	local hppp = entry.hppp
	
	return amount * hppp, healthmax * hppp, true
end

-- DEBUG
if Debug then
	local means = {}
	local function get_statistical_values (samples)
		for i = 1, #samples do
			local amount, percent, heals, damages = unpack(samples[i])
			if 
				(percent > 0 and amount ~= 0 and damages > heals)
			or
				(percent < 0 and amount ~= 100 and heals > damages)
			then
				means[#means + 1] = percent / (damages - heals)
			end
		end
		if not #means then return end
		
		local mean = 0
		for i = 1, #means do
			mean = mean + means[i]
		end
		mean = mean / #means
		local variance = 0
		for i = 1, #means do
			local diff = means[i] - mean
			variance = variance + diff * diff
		end
		local dev = math_sqrt(variance)
		
		return #means, mean, dev
	end
	local debug_frame
	local function get_unit_frame(unit)
		if not debug_frame then
			debug_frame = CreateFrame("Frame", nil, UIParent)
			debug_frame:SetPoint("CENTER")
			debug_frame:SetWidth(120)
			debug_frame:SetHeight(42)
			debug_frame:SetFrameStrata("DIALOG")
			debug_frame:EnableMouse(true)
			debug_frame:SetToplevel(true)
			debug_frame:SetMovable(true)
			debug_frame:SetBackdrop({
				bgFile = "Interface\\TutorialFrame\\TutorialFrameBackground",
				edgeFile = "Interface\\Tooltips\\UI-Tooltip-Border",
				tile = true,
				edgeSize = 16,
				tileSize = 32,
				insets = {
					left = 5,
					right = 5,
					top = 5,
					bottom = 5
				}
			})
			debug_frame:SetScript("OnMouseDown", function(self) self:StartMoving() end)
			debug_frame:SetScript("OnMouseUp", function(self) self:StopMovingOrSizing() end)
			
			local fs = debug_frame:CreateFontString(nil, "OVERLAY")
			fs:SetString("MobHealthDebug")
			fs:SetWidth(120)
			fs:SetHeight(14)
			fs:SetJustifyH("CENTER")
		end
		local frame = debug_frame[unit]
		if frame then return frame end
		frame = CreateFrame("Frame", nil, debug_frame)
		frame:SetPoint("TOPLEFT", 0, unit == "focus" and 28 or 14)
		frame:SetWidth(120)
		frame:SetHeight(12)
		
		local fs = frame:CreateFontString(nil, "OVERLAY")
		frame.unit = fs
		fs:SetString(unit)
		fs:SetPoint("TOPLEFT")
		fs:SetWidth(50)
		fs:SetHeight(12)
		
		fs = frame:CreateFontString(nil, "OVERLAY")
		frame.text = fs
		fs:SetWidth(50)
		fs:SetHeight(12)
		fs:SetPoint("TOPLEFT", 60, 0)
		
		debug_frame[unit] = frame
		return frame
	end

	function MobHealth:UpdateDebugInfo(unit, samples)
		local frame = get_unit_frame(unit)
		local count, mean, deviation = get_statistical_values(samples)
		if not count then
			frame.text:SetText("No valid samples")
		else
			frame.text:SetText(("%.2f +/- %.2f (%d/%d)"):format(mean * 100, deviation * 100, count, #samples))
		end
	end

end

-- REGISTRATION
local function activate(self, oldLib, oldDeactivate)
	if oldLib then
		oldLib:UnitUpdate("target")
		oldLib:UnitUpdate("focus")
		for k, v in pairs(oldLib.db) do
			db[k] = v
		end
		if oldDeactivate then
			oldDeactivate(oldLib)
		end
	end
end

AceLibrary:Register(MobHealth, MAJOR_VERSION, MINOR_VERSION, activate)
